# 导入必要的库和模块
import gradio as gr  # 导入Gradio库，用于创建交互式界面
import numpy as np  # 导入NumPy库，用于数值计算
import pickle  # 导入pickle库，用于加载保存的模型
from PIL import Image  # 导入PIL库中的Image模块，用于图像处理

# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
with open('best_knn_model_k6.pkl', 'rb') as f:
    knn_model = pickle.load(f)  # 使用pickle加载KNN模型

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image):
    # 检查image是否为字典，如果是，则提取图像数据
    if isinstance(image, dict) and "image" in image:
        image = image["image"]

    # 确保输入是 NumPy 数组
    if not isinstance(image, np.ndarray):
        image = np.array(image)

    # 将输入的手写图像转换为 PIL 格式的灰度图，并调整大小为8x8
    try:
        image = Image.fromarray(image).convert('L').resize((8, 8))
    except ValueError as e:
        return f"Error processing image: {e}"

    # 将图像数据转换为数组，并将其展平为1D数组
    img_array = np.array(image).reshape(1, -1)

    # 对图像数据进行归一化处理，使其范围为 0 到 16
    img_array_normalized = (img_array / 16).astype(np.uint8)

    # 进行预测
    prediction = knn_model.predict(img_array_normalized)
    return int(prediction[0])

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
interface = gr.Interface(
    fn=predict_digit,  # 调用的函数
    inputs=gr.Sketchpad(type="pil"),  # 手写板输入，类型为 PIL 图像
    outputs=gr.Label(num_top_classes=1),  # 输出是预测的标签
    title="手写数字识别",  # Web 应用标题
    description="在手写板上输入一个数字，模型将返回预测结果。"  # 描述信息
)

# 启动Gradio接口，用户可以通过这个接口进行交互
interface.launch()


